package cn.android666.mediacodecasynchronous

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.os.Handler
import android.os.HandlerThread
import timber.log.Timber
import java.nio.ByteBuffer
import java.util.concurrent.BlockingQueue
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.TimeUnit


@Suppress("PrivatePropertyName")
class H264Encoder(videoWidth: Int, videoHeight: Int) : MediaCodec.Callback() {

    private val yuvBytesQueue: BlockingQueue<ByteArray> = LinkedBlockingQueue(5)
    private val mMediaCodec: MediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
    private val _1K = 1000
    private val _1M = _1K * 1000
    private val ySize = videoWidth * videoHeight // Y分量大小
    private val oneFrameSize = (ySize * 3) shr 1 // 一帧画面大小
    private var index: Int = 0
    private var temp: Byte = 0
    /** 表示已经调用了close()方法 */
    @Volatile // 因为多线程访问这个变量，所以加上这个注解
    private var calledCloseMethod = false
    private val h264Saver: H264Saver by lazy { H264Saver() }
    private var frameLossCount: Int = 0
    private var isPutEmptyArray = false
    private var mHandlerThread: HandlerThread? = null

    init {
        val mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, videoWidth, videoHeight)
        mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible)
        mediaFormat.setInteger(MediaFormat.KEY_BITRATE_MODE, MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_VBR) // 设置码率为动态码率，默认也是动态码率
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, getBitrate(videoWidth, videoHeight)) // 码率（即比特率）, 官方Demo这里的1000即1kbps，1000_000即1mpbs
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 25)      // 帧速（25帧/秒）
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 2) // I帧间隔（1帧/2秒），因为帧 1秒出25帧，2秒就出50帧，所以I帧间隔为2的话就是每50帧出一个关键帧

        // 第二个参数用于显示解码器的视频内容，第三个参数为编解码器的解密参数，第四个参数为指定为编码器
        mMediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        val handlerThread = HandlerThread("H264EncoderThread").apply { start() }
        mHandlerThread = handlerThread
        val handler = Handler(handlerThread.looper)
        mMediaCodec.setCallback(this, handler) // 传入一个子线程的Handler，以便回调函数可以运行在子线程，如果不传默认运行在主线程
        mMediaCodec.start()
        Timber.i("实际使用的H264编码器：${mMediaCodec.codecInfo.name}")
    }

    override fun onInputBufferAvailable(codec: MediaCodec, index: Int) {
        if (isPutEmptyArray) return // 已经放入了空数组（用于结束编码），则不处理，因为多线程，所以有可能释放的时候还有数据扔进来编码的
        val inputBuffer: ByteBuffer = codec.getInputBuffer(index) ?: return
        try {
            while (true) {
                var yuvBytes = yuvBytesQueue.poll(30L, TimeUnit.MILLISECONDS)
                if (yuvBytes == null && !calledCloseMethod) {
                    continue // 如果已经超时了从队列中取不到数据，并且没有调用close函数，则继续再从队列中再取数据
                }

                val flags = if (calledCloseMethod) {
                    isPutEmptyArray = true
                    Timber.i("已放入空数组")
                    // 如果已经调用了关闭函数，则使用结束flag标志，并放入一个空数组
                    yuvBytes = ByteArray(0)
                    MediaCodec.BUFFER_FLAG_END_OF_STREAM
                } else {
                    nv21ToNv12(yuvBytes)
                    0
                }

                inputBuffer.put(yuvBytes, 0, yuvBytes.size)
                codec.queueInputBuffer(index, 0, yuvBytes.size, System.nanoTime() / 1000, flags)
                break
            }
        } catch (e: Exception) {
            Timber.e(e,"把YUV编码为H264时出现异常")
        }
    }

    override fun onOutputBufferAvailable(codec: MediaCodec, index: Int, info: MediaCodec.BufferInfo) {
        val outputBuffer: ByteBuffer = codec.getOutputBuffer(index) ?: return
        when (info.flags) {
            MediaCodec.BUFFER_FLAG_CODEC_CONFIG, MediaCodec.BUFFER_FLAG_KEY_FRAME, 0 -> { // 配置帧、关键帧、普通帧
                h264Saver.write(outputBuffer)
                mMediaCodec.releaseOutputBuffer(index, false)
            }
            MediaCodec.BUFFER_FLAG_END_OF_STREAM -> {
                Timber.i("已经到达流的终点了")
                mMediaCodec.releaseOutputBuffer(index, false)
                releaseMediaCodec()
            }
        }
    }

    override fun onError(codec: MediaCodec, e: MediaCodec.CodecException) {
        Timber.e(e, "H264编码器出现异常")
    }

    override fun onOutputFormatChanged(codec: MediaCodec, format: MediaFormat) {

    }

    fun addYUVBytes(yuvBytes: ByteArray) {
        if (calledCloseMethod) {
            return
        }

        val offer = yuvBytesQueue.offer(yuvBytes)
        if (!offer) {
            Timber.i("丢帧：${++frameLossCount}帧")
        }
    }

    private fun getBitrate(videoWidth: Int, videoHeight: Int): Int = when {
        (videoWidth == 1920 && videoHeight == 1080) || (videoWidth == 1080 && videoHeight == 1920) -> _1M * 3 shr 1
        (videoWidth == 1280 && videoHeight == 720) || (videoWidth == 720 && videoHeight == 1280) -> _1M
        (videoWidth == 640 && videoHeight == 480) || (videoWidth == 480 && videoHeight == 640) -> _1K * 500
        (videoWidth == 352 && videoHeight == 288) || (videoWidth == 288 && videoHeight == 352) -> _1K * 300
        else -> _1M * 1
    }

    private fun nv21ToNv12(yuvBytes: ByteArray) {
        index = ySize
        while (index < oneFrameSize) {
            temp = yuvBytes[index]
            yuvBytes[index] = yuvBytes[index + 1]
            yuvBytes[index + 1] = temp
            index += 2
        }
    }

    /** 关闭编码器 */
    fun close() {
        Timber.i("close")
        if (calledCloseMethod) return // 预防关闭函数被调用两次
        calledCloseMethod = true
        //yuvBytesQueue.offer(ByteArray(0))不可取，万一队列满了，则无法存进去
    }

    private fun releaseMediaCodec() {
        Timber.i("releaseMediaCodec()")
        mHandlerThread?.quitSafely()
        mHandlerThread = null

        try {
            mMediaCodec.stop()
        } catch (e: Exception) {
            Timber.e(e,"别慌，正常停止${H264Encoder::class.java.simpleName}时出现的异常！")
        }

        try {
            mMediaCodec.release()
        } catch (e: Exception) {
            Timber.e(e,"别慌，正常释放${H264Encoder::class.java.simpleName}时出现的异常！")
        }

        h264Saver.close()
    }

}
